import os
import json

# Directory containing the JSON files
input_dir = r"C:\Users\aikedaer\Desktop\电科院人工智能比赛\traindataset\gt_labels_json"
output_dir = r"C:\Users\aikedaer\Desktop\电科院人工智能比赛\traindataset\gt_labels_txt"

# Ensure output directory exists
os.makedirs(output_dir, exist_ok=True)

# The label mapping dictionary
gt_label_map = {
    "011_zywz_gkzyfgz_gczywzqsyaqd": "1",
    "011_zywz_gkzyfgz_gczywzqsyaqd_fyb": "1",
    "011_wz_xwwz_kyhxcaqwl_sk": "4",
    "011_wz_xwwz_kyhxcaqwl_xz": "4",
    "011_aqfx_fgz_tszywrft_r": "2",
    "011_wz_xwwz_wzqzz_xb": "3",
    "011_zywz_xczyaqjs_zyxckdwzghwfhcs_fyb": "0",
    "011_zywz_xczyaqjs_zyxckdwzghwfhcs": "0",
    "011_wz_zzwz_zyxckdwzg": "0",
    "gczywzqpdaqd": "1",
    "kyhxcaqwl": "4", 
    "tszywrft": "2",
    "wgfzz": "3",
    "zyxckdwzg": "0"
}

# Function to convert bounding box
def convert_bbox(points, img_width, img_height):
    x_min, y_min = points[0]
    x_max, y_max = points[1]
    
    x_center = (x_min + x_max) / 2.0 / img_width
    y_center = (y_min + y_max) / 2.0 / img_height
    width = (x_max - x_min) / img_width
    height = (y_max - y_min) / img_height
    
    return x_center, y_center, width, height

# Process each JSON file in the directory
for filename in os.listdir(input_dir):
    if filename.endswith('.json'):
        file_path = os.path.join(input_dir, filename)
        
        with open(file_path, 'r', encoding="utf-8") as f:
            data = json.load(f)
        
        image_width = data['imageWidth']
        image_height = data['imageHeight']
        
        yolo_annotations = []
        for shape in data['shapes']:
            label = shape['label']
            points = shape['points']
            
            if label in gt_label_map:
                label_id = gt_label_map[label]
                x_center, y_center, width, height = convert_bbox(points, image_width, image_height)
                yolo_annotations.append(f"{label_id} {x_center} {y_center} {width} {height}")
        
        # Write YOLO annotations to a .txt file
        output_filename = os.path.splitext(filename)[0] + '.txt'
        output_file_path = os.path.join(output_dir, output_filename)
        
        with open(output_file_path, 'w') as out_f:
            for annotation in yolo_annotations:
                out_f.write(f"{annotation}\n")
